#include "ScreenCapture.h"
#include <iostream>
#include <conio.h>

// Function to check if running with elevated privileges
bool IsRunningAsAdministrator() {
    HANDLE hToken = NULL;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken)) {
        TOKEN_ELEVATION elevation;
        DWORD dwSize;
        if (GetTokenInformation(hToken, TokenElevation, &elevation, sizeof(elevation), &dwSize)) {
            CloseHandle(hToken);
            return elevation.TokenIsElevated != 0;
        }
        CloseHandle(hToken);
    }
    return false;
}

// Function to check if running as SYSTEM
bool IsRunningAsSystem() {
    HANDLE hToken = NULL;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken)) {
        DWORD dwSize = 0;
        GetTokenInformation(hToken, TokenUser, NULL, 0, &dwSize);

        if (dwSize > 0) {
            PTOKEN_USER pTokenUser = (PTOKEN_USER)malloc(dwSize);
            if (GetTokenInformation(hToken, TokenUser, pTokenUser, dwSize, &dwSize)) {
                char userName[256];
                char domainName[256];
                DWORD userNameSize = sizeof(userName);
                DWORD domainNameSize = sizeof(domainName);
                SID_NAME_USE snu;

                if (LookupAccountSidA(NULL, pTokenUser->User.Sid, userName, &userNameSize,
                                     domainName, &domainNameSize, &snu)) {
                    bool isSystem = (strcmp(userName, "SYSTEM") == 0);
                    free(pTokenUser);
                    CloseHandle(hToken);
                    return isSystem;
                }
            }
            free(pTokenUser);
        }
        CloseHandle(hToken);
    }
    return false;
}

int main() {
    std::cout << "ScreenCapture Service Test Program" << std::endl;
    std::cout << "==================================" << std::endl;
    std::cout << std::endl;

    // Check privileges
    if (IsRunningAsSystem()) {
        std::cout << "✓ Running as SYSTEM - Full secure desktop access available" << std::endl;
    } else if (IsRunningAsAdministrator()) {
        std::cout << "✓ Running as Administrator - Enhanced privileges available" << std::endl;
    } else {
        std::cout << "⚠ Running as regular user - Secure desktop capture may be limited" << std::endl;
        std::cout << "  For full secure desktop access, run as Administrator or SYSTEM service" << std::endl;
    }
    std::cout << std::endl;

    ScreenCapture capture;
    CaptureSettings settings;

    settings.method = CaptureMethod::AUTO;
    settings.outputPath = "C:\\ScreenCapture";
    settings.captureInterval = 3000;  // 3 seconds for testing
    settings.enableDesktopMonitoring = true;
    settings.captureSecureDesktops = true;

    std::cout << "Initializing enhanced screen capture..." << std::endl;
    CaptureResult result = capture.Initialize(settings);

    if (result != CaptureResult::SUCCESS) {
        std::cout << "Failed to initialize: " << capture.GetLastError() << std::endl;
        std::cout << "Press any key to exit..." << std::endl;
        _getch();
        return 1;
    }

    std::cout << "✓ Initialization successful!" << std::endl;
    std::cout << "Output directory: " << settings.outputPath << std::endl;
    std::cout << std::endl;

    std::cout << "Instructions:" << std::endl;
    std::cout << "1. Press 's' to start continuous capture" << std::endl;
    std::cout << "2. Lock your screen (Win+L) to test secure desktop capture" << std::endl;
    std::cout << "3. Check the output folder for screenshots" << std::endl;
    std::cout << "4. Press 'q' to quit" << std::endl;
    std::cout << std::endl;

    bool running = true;
    bool capturing = false;

    while (running) {
        std::cout << "Command (s=start, t=stop, c=single capture, q=quit): ";
        char cmd = _getch();
        std::cout << cmd << std::endl;

        switch (cmd) {
            case 's':
            case 'S':
                if (!capturing) {
                    std::cout << "Starting continuous capture..." << std::endl;
                    result = capture.StartCapture();
                    if (result == CaptureResult::SUCCESS) {
                        std::cout << "✓ Capture started! Lock screen to test secure desktop capture." << std::endl;
                        capturing = true;
                    } else {
                        std::cout << "Failed to start capture: " << capture.GetLastError() << std::endl;
                    }
                } else {
                    std::cout << "Capture already running!" << std::endl;
                }
                break;

            case 't':
            case 'T':
                if (capturing) {
                    std::cout << "Stopping capture..." << std::endl;
                    result = capture.StopCapture();
                    if (result == CaptureResult::SUCCESS) {
                        std::cout << "✓ Capture stopped." << std::endl;
                        capturing = false;
                    } else {
                        std::cout << "Failed to stop capture." << std::endl;
                    }
                } else {
                    std::cout << "Capture not running!" << std::endl;
                }
                break;

            case 'c':
            case 'C':
                std::cout << "Taking single screenshot..." << std::endl;
                result = capture.CaptureFrame();
                if (result == CaptureResult::SUCCESS) {
                    std::cout << "✓ Screenshot saved!" << std::endl;
                } else {
                    std::cout << "Screenshot failed: " << capture.GetLastError() << std::endl;
                }
                break;

            case 'q':
            case 'Q':
                std::cout << "Exiting..." << std::endl;
                running = false;
                break;

            default:
                std::cout << "Invalid command." << std::endl;
                break;
        }
        std::cout << std::endl;
    }

    if (capturing) {
        std::cout << "Stopping capture before exit..." << std::endl;
        capture.StopCapture();
    }

    std::cout << "Program ended. Check C:\\ScreenCapture for screenshots." << std::endl;
    std::cout << "Press any key to exit..." << std::endl;
    _getch();

    return 0;
}